package com.zhigh.project.security.jwt.web.authentication;

import com.zhigh.project.security.jwt.authentication.DefaultJWT;
import com.zhigh.project.security.jwt.authentication.JWTAuthentication;
import com.zhigh.project.security.jwt.authentication.JWTMessageSource;
import com.zhigh.project.security.jwt.authentication.event.JWTAuthenticationFailureEvent;
import com.zhigh.project.security.jwt.authentication.principal.DefaultJWTPrincipal;
import com.zhigh.project.security.jwt.web.token.TokenExtractor;
import com.zhigh.project.tools.Assertions;
import org.apache.commons.lang3.StringUtils;
import org.springframework.context.ApplicationEventPublisher;
import org.springframework.context.ApplicationEventPublisherAware;
import org.springframework.context.MessageSource;
import org.springframework.context.MessageSourceAware;
import org.springframework.context.support.MessageSourceAccessor;
import org.springframework.core.log.LogMessage;
import org.springframework.lang.NonNull;
import org.springframework.security.authentication.AuthenticationManager;
import org.springframework.security.authentication.InternalAuthenticationServiceException;
import org.springframework.security.authentication.event.InteractiveAuthenticationSuccessEvent;
import org.springframework.security.core.Authentication;
import org.springframework.security.core.AuthenticationException;
import org.springframework.security.core.context.SecurityContext;
import org.springframework.security.core.context.SecurityContextHolder;
import org.springframework.security.web.authentication.AuthenticationFailureHandler;
import org.springframework.security.web.authentication.AuthenticationSuccessHandler;
import org.springframework.security.web.util.matcher.RequestHeaderRequestMatcher;
import org.springframework.security.web.util.matcher.RequestMatcher;
import org.springframework.web.filter.OncePerRequestFilter;

import javax.servlet.FilterChain;
import javax.servlet.ServletException;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.io.IOException;

public class JWTAuthenticationFilter extends OncePerRequestFilter implements MessageSourceAware, ApplicationEventPublisherAware {

    private final static String DEFAULT_AUTHENTICATION_HEADER_NAME = "Authorization";

    protected MessageSourceAccessor messages = JWTMessageSource.getAccessor();

    private TokenExtractor tokenExtractor;

    private RequestMatcher requiresAuthenticationRequestMatcher;

    private AuthenticationManager authenticationManager;

    private boolean continueChainAfterSuccessfulAuthentication = true;

    private ApplicationEventPublisher eventPublisher;

    private AuthenticationSuccessHandler successHandler = (request, response, authentication) -> {
    };

    private AuthenticationFailureHandler failureHandler = (request, response, authentication) -> {
    };

    public JWTAuthenticationFilter(AuthenticationManager authenticationManager) {
        this(DEFAULT_AUTHENTICATION_HEADER_NAME, authenticationManager);
    }

    public JWTAuthenticationFilter(String tokenHoldHeaderName, AuthenticationManager authenticationManager) {
        Assertions.assertNonnull(tokenHoldHeaderName, "tokenHoldHeaderName must non-blank.");
        this.requiresAuthenticationRequestMatcher = new RequestHeaderRequestMatcher(tokenHoldHeaderName);
        this.authenticationManager = authenticationManager;
        this.tokenExtractor = new HeaderTokenExtractor(tokenHoldHeaderName);
    }

    public JWTAuthenticationFilter(RequestMatcher requiresAuthenticationRequestMatcher, TokenExtractor tokenExtractor, AuthenticationManager authenticationManager) {
        this.requiresAuthenticationRequestMatcher = requiresAuthenticationRequestMatcher;
        this.tokenExtractor = tokenExtractor;
        this.authenticationManager = authenticationManager;
    }

    @Override
    protected void doFilterInternal(@NonNull HttpServletRequest request, @NonNull HttpServletResponse response, @NonNull FilterChain filterChain) throws ServletException, IOException {
        if (!requiresAuthenticationRequestMatcher.matches(request)) {
            filterChain.doFilter(request, response);
            return;
        }

        Authentication authResult = null;
        try {
            String token = tokenExtractor.extract(request);

            DefaultJWT jwt = new DefaultJWT(token);

            authResult = authenticationManager.authenticate(new JWTAuthentication(new DefaultJWTPrincipal(jwt)));

            successfulAuthentication(request, response, filterChain, authResult);

            if (this.continueChainAfterSuccessfulAuthentication) {
                filterChain.doFilter(request, response);
            }

        } catch (InternalAuthenticationServiceException failed) {
            this.logger.error("An internal error occurred while trying to authenticate the user.", failed);
            unsuccessfulAuthentication(request, response, authResult, failed);
        } catch (AuthenticationException ex) {
            this.logger.error("An error occurred while trying to authenticate the user.", ex);
            unsuccessfulAuthentication(request, response, authResult, ex);
        }
    }

    protected void successfulAuthentication(HttpServletRequest request, HttpServletResponse response, FilterChain filterChain, Authentication authenticationResult) throws ServletException, IOException {
        SecurityContext context = SecurityContextHolder.createEmptyContext();
        context.setAuthentication(authenticationResult);
        SecurityContextHolder.setContext(context);
        if (this.logger.isDebugEnabled()) {
            this.logger.debug(LogMessage.format("Set SecurityContextHolder to %s", authenticationResult));
        }

        if (eventPublisher != null) {
            eventPublisher.publishEvent(new InteractiveAuthenticationSuccessEvent(authenticationResult, this.getClass()));
        }

        this.successHandler.onAuthenticationSuccess(request, response, authenticationResult);

    }

    protected void unsuccessfulAuthentication(HttpServletRequest request, HttpServletResponse response, Authentication authenticationResult, AuthenticationException failed) throws ServletException, IOException {
        SecurityContextHolder.clearContext();
        this.logger.trace("Failed to process authentication request", failed);
        this.logger.trace("Cleared SecurityContextHolder");
        this.logger.trace("Handling authentication failure");

        if (eventPublisher != null) {
            eventPublisher.publishEvent(new JWTAuthenticationFailureEvent(authenticationResult, failed));
        }

        failureHandler.onAuthenticationFailure(request, response, failed);
    }


    @Override
    public void setMessageSource(@NonNull MessageSource messageSource) {
        this.messages = new MessageSourceAccessor(messageSource);
    }

    public void setTokenExtractor(TokenExtractor tokenExtractor) {
        this.tokenExtractor = tokenExtractor;
    }

    public void setRequiresAuthenticationRequestMatcher(RequestMatcher requiresAuthenticationRequestMatcher) {
        this.requiresAuthenticationRequestMatcher = requiresAuthenticationRequestMatcher;
    }

    @Override
    public void afterPropertiesSet() throws ServletException {
        super.afterPropertiesSet();
        Assertions.assertNonnull(authenticationManager, "AuthenticationManager is required to set.");
        Assertions.assertNonnull(requiresAuthenticationRequestMatcher, "requiresAuthenticationRequestMatcher is required to set.");
        Assertions.assertNonnull(messages, "Messages is required to set.");
        Assertions.assertNonnull(tokenExtractor, "tokenExtractor is required to set.");
    }

    public void setSuccessHandler(AuthenticationSuccessHandler successHandler) {
        this.successHandler = successHandler;
    }

    public void setFailureHandler(AuthenticationFailureHandler failureHandler) {
        this.failureHandler = failureHandler;
    }

    @Override
    public void setApplicationEventPublisher(@NonNull ApplicationEventPublisher eventPublisher) {
        this.eventPublisher = eventPublisher;
    }

    public void setAuthenticationManager(AuthenticationManager authenticationManager) {
        this.authenticationManager = authenticationManager;
    }

    public void setContinueChainAfterSuccessfulAuthentication(boolean continueChainAfterSuccessfulAuthentication) {
        this.continueChainAfterSuccessfulAuthentication = continueChainAfterSuccessfulAuthentication;
    }

    private class HeaderTokenExtractor implements TokenExtractor {

        private final String exceptHeaderName;

        public HeaderTokenExtractor(String exceptHeaderName) {
            this.exceptHeaderName = exceptHeaderName;
        }

        @Override
        public String extract(HttpServletRequest request) {

            String headerValue = request.getHeader(exceptHeaderName);
            if (StringUtils.isBlank(headerValue)) return "";
            String token = headerValue.substring(headerValue.indexOf(" ") + 1);
            logger.trace("extract token: " + token);
            return token;
        }
    }
}
